package edu.colorado.cs.correll.cudroplet.shaperecognition;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.xml.bind.JAXB;

import com.vividsolutions.jts.algorithm.match.AreaSimilarityMeasure;
import com.vividsolutions.jts.algorithm.match.HausdorffSimilarityMeasure;
import com.vividsolutions.jts.algorithm.match.SimilarityMeasure;
import com.vividsolutions.jts.geom.Coordinate;
import com.vividsolutions.jts.geom.Geometry;
import com.vividsolutions.jts.geom.GeometryFactory;

import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.ListTestRun;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.ListShapePosition;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.ListTestType;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.Match;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.TestRun;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.ShapePosition;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.ShapeRecognitionRuns;
import edu.colorado.cs.correll.cudroplet.shaperecognition.xml.TestType;
import processing.core.PApplet;
import processing.core.PShape;
import processing.core.PVector;

@SuppressWarnings("serial")
public class ShapeRecognition extends PApplet {
	private static final int REPEATING_LINES_WEIGHT = 10;

	private static final float SIMILARITY_STROKE_WEIGHT = 3;

	// size of the robot body
	private static final int ROBOT_SIZE = 3;

	// the radius to show the leader
	private static final int LEADER_DIAMETER = 5;

	// the weight of the line that connects neighbors
	private static final int NEIGHBOR_LINE_SIZE = 1;

	// the weight of the simple polygon line
	private static final int ALPHA_SHAPE_LINE_STROKE_WEIGHT = 3;

	private static final int POINT_WEIGHT = 10;

	private static final int LABEL_OFFSET = 5;

	private static final boolean TEST = false;

	private static final int MATCH_LINE_OFFSET = 5;

	private static final int LEGEND_LINE_HEIGHT = 12;

	private static final int ROBOT_DOT_SIZE = 2;

	private static final int ALPHA_SHAPE_CIRCUMCIRCLE_STROKE_WEIGHT = 1;

	private static final int ELLIPSES_OFFSET = 500;

	private static final boolean DEBUG_TRI = false;

	private static final boolean DEBUG_CIRCUMCIRCLES = true;

	private static final int BUFFER = 10;

	private static final int SUBDIVISIONS = 1;

	private static final int COLLECTION_RATE = 1000;

	private static final float FRAME_RATE = 120;
	
	List<Test> tests = new LinkedList<Test>();

	private static final int NUM_COLLECTIONS = 20;

	private static final double HAUSDORFF_THRESHOLD = 0.95;
	private static final double AREA_THRESHOLD = 0.9;

	private final int ALPHA_RADIUS_COLOR = color(0, 100, 50);

	private final int EXCLUSION_RADIUS_COLOR = color(50, 100, 50);

	private final int COMM_RADIUS_COLOR = color(100, 0, 50);

	Simulation simulation;

	private AreaSimilarityMeasure areaSimilarityMeasure = new AreaSimilarityMeasure();

	private HausdorffSimilarityMeasure hausdorffSimilarityMeasure = new HausdorffSimilarityMeasure();

	GeometryFactory geometryFactory = new GeometryFactory();
	
	ShapeRecognitionRuns shapeRecognitionRuns = new ShapeRecognitionRuns();

	TestRun run = null;
	
	public ShapeRecognition() {
		Map<SimilarityMeasure, Double> thresholds = new HashMap<SimilarityMeasure, Double>();
		thresholds.put(hausdorffSimilarityMeasure, new Double(HAUSDORFF_THRESHOLD));
		thresholds.put(areaSimilarityMeasure, new Double(AREA_THRESHOLD));
		
		SimilarityMeasure[] similarityMeasures = {hausdorffSimilarityMeasure, areaSimilarityMeasure};
		ShapeType[] shapeTypes = {ShapeType.CONVEX_HULL, ShapeType.UNION_OF_BALLS, ShapeType.ALPHA_SHAPE};
		
		ColorGenerator colorGenerator = new ColorGenerator(similarityMeasures.length * shapeTypes.length);
		int i = 0;
		for(SimilarityMeasure similarityMeasure: similarityMeasures) {
			for(ShapeType shapeType: shapeTypes) {
				tests.add(new Test(i, similarityMeasure, shapeType, thresholds.get(similarityMeasure), colorGenerator.getColor(i)));
				i++;
			}
		}
	}

	public void setup() {
		size(1000, 1000);
		frameRate(FRAME_RATE);
	}

	public void draw() {
		
		if((frameCount % COLLECTION_RATE) == 1) {
			if (TEST) {
				simulation = new TestSimulation(this);
			} else {
				simulation = new SimpleSimulation(this);
			}
			run = new TestRun();
			run.number = frameCount / COLLECTION_RATE;
		}
		if(frameCount == 1) {
			shapeRecognitionRuns.swarmSize = Swarm.SWARM_SIZE;
			shapeRecognitionRuns.stepRate = simulation.swarm.robotSettings.step_rate;
			shapeRecognitionRuns.commRadius = simulation.swarm.robotSettings.comm_radius;
			shapeRecognitionRuns.exclusionRadius = simulation.swarm.robotSettings.exclusion_radius;
			shapeRecognitionRuns.train = simulation.train;
			shapeRecognitionRuns.test = simulation.test;
			shapeRecognitionRuns.shapeDim = EnvironmentMap.SHAPE_DIM;
			shapeRecognitionRuns.listTestRun = new ListTestRun();
			shapeRecognitionRuns.listShapePosition = new ListShapePosition();
			shapeRecognitionRuns.listTestType = new ListTestType();
			
			for(PlacedShape placedShape: simulation.testPlacedShapes) {
				ShapePosition shapePosition = new ShapePosition();
				shapePosition.shapeName = placedShape.name;
				shapePosition.x = placedShape.position.x;
				shapeRecognitionRuns.listShapePosition.shapePositionList.add(shapePosition);
			}
			
			for(Test test: tests) {
				TestType testType = new TestType();
				testType.id = test.id;
				testType.similarityMeasure = Test.getSimilarityMeasureName(test.similarityMeasure);
				testType.shapeType = Test.getShapeTypeName(test.shapeType);
				testType.threshold = test.threshold;
			}
		}
		
		simulation.draw();
		if((frameCount % COLLECTION_RATE) == 0) {
			drawAlphaShapes(true);
		    matchShapes(run);
		    drawLegend(100, 100, simulation.swarm.robotSettings);
		    shapeRecognitionRuns.listTestRun.testRunList.add(run);
		    save("result" + String.format("%04d", frameCount/COLLECTION_RATE) + ".png");
		}
		if(frameCount == COLLECTION_RATE * NUM_COLLECTIONS) {
			JAXB.marshal(shapeRecognitionRuns, createOutput("results.xml"));
			exit();
		}
	}
	

	public void drawConvexHulls(boolean bDrawLeader) {
		for (Subgraph subgraph : simulation.swarm.subgraphs) {
			Geometry convexHull = subgraph.getMultiPoint().convexHull();
			drawGeometry(subgraph.max.id, ALPHA_SHAPE_LINE_STROKE_WEIGHT,
					convexHull);
			if (bDrawLeader) {
				drawLeaderLabel(subgraph.max.id, convexHull);
			}
		}
	}

	private void drawGeometry(int id, int stroke_weight, Geometry geometry) {
		int c1 = simulation.swarm.getColor(id);
		int c2 = color(100);
		for (Line line : GeometryUtil.getLines(geometry)) {
			biline(line.start.x, line.start.y, line.end.x - line.start.x,
					line.end.y - line.start.y, stroke_weight, c1, c2);
		}
	}

	public void drawSubgraphs() {
		for (Subgraph s : simulation.swarm.subgraphs) {
			drawSubgraph(s.max.id, POINT_WEIGHT, s);
		}
	}

	void drawSubgraph(int maxId, int weight, Subgraph subgraph) {
		int c = simulation.swarm.getColor(maxId);
		fillFromInt(c);

		for (Robot r : subgraph.subgraphSet) {
			strokeFromInt(c);
			text(r.id, r.x + LABEL_OFFSET, r.y);
			noStroke();
			ellipse(r.x, r.y, weight, weight);
		}
	}

	public void drawAlphaShapes(boolean bDrawLeader) {
		for (Subgraph subgraph : simulation.swarm.subgraphs) {
			AlphaShapeCollection alphaShapeCollection = subgraph
					.getAlphaShapeCollection();
			for (AlphaShape alphaShape : alphaShapeCollection.alphaShapes) {
				if (DEBUG_CIRCUMCIRCLES) {
					drawCircumcircles(alphaShape.max.id,
							ALPHA_SHAPE_LINE_STROKE_WEIGHT,
							ALPHA_SHAPE_CIRCUMCIRCLE_STROKE_WEIGHT,
							ROBOT_DOT_SIZE, alphaShape.circumcircles);
					
				}
				drawGeometry(alphaShape.max.id, ALPHA_SHAPE_LINE_STROKE_WEIGHT, alphaShape.getGeometry().buffer(BUFFER, SUBDIVISIONS)); 
			
				if (bDrawLeader) {
					Coordinate[] coordinates = alphaShape.getCoordinates();
					drawLeaderLabel(alphaShape.max.id,
							geometryFactory.createMultiPoint(coordinates));
				}
			}
		}
	}

	private void drawCircumcircles(int id, int line_stroke_weight,
			int circumcircle_stroke_weight, int body_size,
			Collection<Circumcircle> circumcircles) {
		int c = simulation.swarm.getColor(id);

		for (Circumcircle circumcircle : circumcircles) {
			Coordinate center = circumcircle.getCenter();
			Coordinate coordA = circumcircle.a.getCoordinate();
			Coordinate coordB = circumcircle.b.getCoordinate();
			drawCircumcircle(center, circumcircle.diameter / 2,
					circumcircle_stroke_weight, c);
			bilineFromCoords(coordA, coordB, line_stroke_weight, c, color(100));
			if (DEBUG_TRI) {
				Coordinate[] tri = { center, coordA, coordB };
				for (int i = 0; i < tri.length; i++) {
					drawCoordinate(tri[i], body_size, c);
					for (int j = i; j < tri.length; j++) {
						bilineFromCoords(tri[i], tri[j], line_stroke_weight, c,
								color(100));
					}
				}
			}
		}
	}

	private void bilineFromCoords(Coordinate a, Coordinate b,
			int stroke_weight, int c1, int c2) {
		biline((float) a.x, (float) a.y, (float) (b.x - a.x),
				(float) (b.y - a.y), stroke_weight, c1, c2);

	}

	private void drawCoordinate(Coordinate coord, int body_size, int color) {
		fillFromInt(color);
		ellipse((float) coord.x, (float) coord.y, body_size, body_size);
	}

	private void drawCircumcircle(Coordinate center, float radius,
			int stroke_weight, int color) {

		strokeWeight(stroke_weight);
		strokeFromInt(color);
		noFill();
		ellipse((float) center.x, (float) center.y, radius * 2, radius * 2);

	}

	private void drawLeaderLabel(int maxId, Geometry geom) {
		Coordinate centroid = geom.getCentroid().getCoordinate();
		fill(0);
		text(maxId, (float) centroid.x, (float) centroid.y);
	}

	/**
	 * draw the exclusion radius, used for picking
	 */
	void drawExclusion(Robot r) {
		if (r.state == RobotState.FOUND) {
			fill(255);
			noStroke();
			ellipse(r.x, r.y,
					simulation.swarm.robotSettings.exclusion_radius * 2,
					simulation.swarm.robotSettings.exclusion_radius * 2);
		}
	}

	/**
	 * Indicate the leader
	 */
	void drawLeader(Robot r) {
		if (r.leader != null) {
			int c = simulation.swarm.getColor(r.leader.id);
			fillFromInt(c);
			noStroke();
			ellipse(r.x, r.y, LEADER_DIAMETER, LEADER_DIAMETER);
		}
	}

	/**
	 * draw the body and comm radius, not used for picking
	 * 
	 * @param drawNumber
	 */
	void drawBody(Robot r, boolean simple, boolean drawNumber) {
		if (!simple) {
			int head = 0;
			if (r.state == RobotState.SEARCH) {
				head = color(255, 0, 255);
			} else if (r.state == RobotState.FOUND) {
				head = color(255, 255, 0);
			}
			int tail = blendColor(head, color(100, 100, 100), BLEND);
			biline(r.x, r.y, 2 * ROBOT_SIZE * cos(r.theta), 2 * ROBOT_SIZE
					* sin(r.theta), ROBOT_SIZE, head, tail);
		}
		if (r.state == RobotState.FOUND) {
			if (drawNumber) {
				fill(0);
				text(r.id, r.x + 20, r.y);
			}
			if (!simple) {
				noFill();
				strokeWeight(1);
				stroke(255);
			} else {
				noFill();
				strokeWeight(1);
				stroke(255);
				noStroke();
			}
			ellipse(r.x, r.y, simulation.swarm.robotSettings.comm_radius * 2,
					simulation.swarm.robotSettings.comm_radius * 2);
			fillFromInt(color(100));
			//ellipse(r.x, r.y, ROBOT_SIZE, ROBOT_SIZE);

			if (!simple) {
				for (Robot r2 : simulation.swarm.neighbors[r.id]) {
					biline(r.x, r.y, r2.x - r.x, r2.y - r.y,
							NEIGHBOR_LINE_SIZE, color(100), color(200));
				}
			}
		}
	}

	// draw a line that has a head and a tail
	public void biline(float x, float y, float len_x, float len_y, int weight,
			int head, int tail) {
		strokeWeight(weight);
		strokeFromInt(head);
		line(x, y, x + len_x / 2, y + len_y / 2);
		strokeFromInt(tail);
		line(x + len_x / 2, y + len_y / 2, x + len_x, y + len_y);
	}

	void printOnce(String s) {
		if (frameCount < 2) {
			print(s);
		}
	}

	void printlnOnce(String s) {
		if (frameCount < 2) {
			println(s);
		}
	}

	void myassert(boolean test, String s) {
		if (!test) {
			println(s);
			assert (test);
		}
	}

	void drawOpenChain(int id, int stroke_weight, ArrayList<PVector> vertices) {
		drawShape(id, stroke_weight, vertices, false);
	}

	void drawPolygon(int id, int stroke_weight, ArrayList<PVector> vertices) {
		drawShape(id, stroke_weight, vertices, true);
	}

	void drawShape(int id, int stroke_weight, ArrayList<PVector> vertices,
			boolean closed) {
		int c = simulation.swarm.getColor(id);

		for (int i = 0; i < vertices.size() - (closed ? 0 : 1); i++) {
			PVector curr = vertices.get(i);
			PVector next = vertices.get((i + 1) % vertices.size());
			PVector len = PVector.sub(curr, next);

			biline(next.x, next.y, len.x, len.y, stroke_weight, c, color(100));
		}
	}

	void drawExclusions() {
		for (int i = 0; i < simulation.swarm.robots.length; i++) {
			drawExclusion(simulation.swarm.robots[i]);
		}
	}

	public void drawBodies(boolean simple, boolean drawNumber) {
		for (int i = 0; i < simulation.swarm.robots.length; i++) {
			drawBody(simulation.swarm.robots[i], simple, drawNumber);
		}

	}

	String getLoc() {
		String baseLoc = "/Users/charlesdietrich"; // "/home/common/"
		return baseLoc
				+ "/Dropbox/cu-droplet/simulator/processing/shape-recognition/data/";
	}

	@Override
	public PShape loadShape(String filename) {
		return super.loadShape(getLoc() + filename);
	}

	public void highlightLines(int maxId, Set<Line> repeatingLines) {
		int c = simulation.swarm.getColor(maxId);
		for (Line line : repeatingLines) {
			PVector len = PVector.sub(line.start, line.end);
			biline(line.end.x, line.end.y, len.x, len.y,
					REPEATING_LINES_WEIGHT, c, color(0));
		}

	}

	public void matchShapes(TestRun testRun) {

		Subgraph[] subgraphs = new Subgraph[simulation.swarm.subgraphs.size()];
		simulation.swarm.subgraphs.toArray(subgraphs);
		Geometry[] convexHulls = new Geometry[subgraphs.length];
		AlphaShapeCollection[] alphaShapeCollections = new AlphaShapeCollection[subgraphs.length];
		Geometry[] centeredUnionOfBalls = new Geometry[subgraphs.length];
		Coordinate[] centroids = new Coordinate[subgraphs.length];
		List<AlphaShape> alphaShapes = new LinkedList<AlphaShape>();

		boolean doConvexHull = false;
		boolean doAlphaShape = false;
		boolean doUnionOfBalls = false;
		for(Test test: tests) {
			if(test.shapeType.equals(ShapeType.CONVEX_HULL)) {
				doConvexHull = true;
			}
			
			if(test.shapeType.equals(ShapeType.ALPHA_SHAPE)) {
				doAlphaShape = true;
			}
			if(test.shapeType.equals(ShapeType.UNION_OF_BALLS)) {
				doUnionOfBalls = true;
			}
		}
		for (int i = 0; i < subgraphs.length; i++) {
			Subgraph subgraph = subgraphs[i];
			Geometry centeredMultiPoint = subgraph.getCenteredMultiPoint();
			centroids[i] = subgraph.getMultiPoint().getCentroid()
					.getCoordinate();
			if (doConvexHull) {
				convexHulls[i] = centeredMultiPoint.convexHull();
			}
			if (doAlphaShape) {
				alphaShapeCollections[i] = subgraph.getAlphaShapeCollection();
				alphaShapes.addAll(alphaShapeCollections[i].alphaShapes);
			}
			if (doUnionOfBalls) {
				centeredUnionOfBalls[i] = centeredMultiPoint
						.buffer(BUFFER, SUBDIVISIONS);
			}
		}

		strokeWeight(SIMILARITY_STROKE_WEIGHT);

		for (int i = 0; i < subgraphs.length; i++) {
			for (int j = 0; j < i; j++) {
				for(Test test: tests) {
					if(test.shapeType.equals(ShapeType.CONVEX_HULL) || test.shapeType.equals(ShapeType.UNION_OF_BALLS)) {
						Geometry[] geoms = null;
						if(test.shapeType.equals(ShapeType.CONVEX_HULL)) {
							geoms = convexHulls;
						} else if(test.shapeType.equals(ShapeType.UNION_OF_BALLS)) {
							geoms = centeredUnionOfBalls;
						}
						match(test, geoms, centroids, i, j, testRun);
					}
				}
			}
		}
		if(doAlphaShape) {
			AlphaShape[] alphaShapeArray = new AlphaShape[alphaShapes.size()];
			alphaShapes.toArray(alphaShapeArray);
			Coordinate[] alphaShapeCentroids = new Coordinate[alphaShapes.size()];
			Geometry[] alphaShapeGeoms = new Geometry[alphaShapes.size()];
			for(int i = 0; i < alphaShapeArray.length; i++) {
				Coordinate[] coordinates = alphaShapeArray[i].getCoordinates();
				alphaShapeCentroids[i] = geometryFactory.createMultiPoint(coordinates).getCentroid().getCoordinate();
				alphaShapeGeoms[i] = alphaShapeArray[i].getCenteredGeometry(alphaShapeCentroids[i]).buffer(BUFFER, SUBDIVISIONS);
			}
			
			for (int i = 0; i < alphaShapeArray.length; i++) {
				for (int j = 0; j < i; j++) {
					for(Test test: tests) {
						if(test.shapeType.equals(ShapeType.ALPHA_SHAPE)) {
							match(test, alphaShapeGeoms, alphaShapeCentroids, i, j, testRun);
						}
					}
				}
			}
		}
	}

	private void match(Test test, Geometry[] geoms, Coordinate[] centroids, int i, int j, TestRun testRun) {

		double result = test.similarityMeasure.measure(geoms[i], geoms[j]);
		if (result > test.threshold) {
			strokeFromInt(test.color);
			line((float) centroids[i].x + MATCH_LINE_OFFSET * test.id, (float) centroids[i].y, (float) centroids[j].x
					+ MATCH_LINE_OFFSET * test.id, (float) centroids[j].y);
		}
		
		Match match = new Match();
		match.testId = test.id;
		match.result = result;
		match.startX = centroids[i].x;
		match.startY = centroids[i].y;
		match.endX = centroids[j].x;
		match.endY = centroids[j].y;
		testRun.matchList.add(match);
		
	}

	private void strokeFromInt(int color) {
		stroke(red(color), green(color), blue(color));
	}

	private void fillFromInt(int color) {
		fill(red(color), green(color), blue(color));
	}

	public void drawLegend(float x, float y, RobotSettings robotSettings) {
		int idx = 0;
		for(Test test: tests) {
			fillFromInt(test.color);
			text(Test.getSimilarityMeasureName(test.similarityMeasure) + " " + Test.getShapeTypeName(test.shapeType) + " threshold = " + test.threshold,
					x, y + idx++ * LEGEND_LINE_HEIGHT);
		}
		
		fill(0);
		text("frame " + frameCount, x, y + LEGEND_LINE_HEIGHT * idx++);

		noFill();
		strokeFromInt(COMM_RADIUS_COLOR);
		ellipse(x + ELLIPSES_OFFSET, y, robotSettings.comm_radius * 2,
				robotSettings.comm_radius * 2);
		strokeFromInt(EXCLUSION_RADIUS_COLOR);
		ellipse(x + ELLIPSES_OFFSET, y, robotSettings.exclusion_radius * 2,
				robotSettings.exclusion_radius * 2);
		strokeFromInt(ALPHA_RADIUS_COLOR);
		ellipse(x + ELLIPSES_OFFSET, y, robotSettings.comm_radius,
				robotSettings.comm_radius);

	}
}